import argparse
import m5
from m5.objects import *

# 创建系统
system = System()
system.clk_domain = SrcClockDomain(clock="1GHz")
system.clk_domain.voltage_domain = VoltageDomain()

system.mem_ranges = [AddrRange("512MB")]

# 创建CPU（通过命令行参数选择）
def setup_cpu(cpu_type):
    if cpu_type == "atomic":
        return RiscvAtomicSimpleCPU()
    elif cpu_type == "trace":
        return RiscvTraceCPU(cpu_id=0)
    elif cpu_type == "timing":
        return RiscvTimingSimpleCPU()
    else:
        raise ValueError("Unknown CPU type")

def setup_mem(cpu_type):
    if cpu_type == "atomic":
        return "atomic"
    else:
        return "timing"

# 命令行参数解析
parser = argparse.ArgumentParser()
parser.add_argument("--cpu-type", choices=["atomic", "trace", "timing"], required=True)
args = parser.parse_args()

# 设置CPU和内存
system.cpu = setup_cpu(args.cpu_type)
system.mem_mode = setup_mem(args.cpu_type)               # 内存模式需与CPU匹配

system.membus = SystemXBar()
system.cpu.icache_port = system.membus.cpu_side_ports[0]
system.cpu.dcache_port = system.membus.cpu_side_ports[1]
system.cpu.createInterruptController()

# 设置内存
system.mem_ctrl = MemCtrl()
system.mem_ctrl.dram = DDR3_1600_8x8()
system.mem_ctrl.port = system.membus.mem_side_ports[0]

# 设置进程
thispath = os.path.dirname(os.path.realpath(__file__))
binary = os.path.join(
    thispath,
    "riscv-binaries",
    "hello-world",
)

system.workload = SEWorkload.init_compatible(binary)

process = Process()
process.cmd = [binary]                   # 测试程序路径
system.cpu.workload = process
system.cpu.createThreads()

# 启动仿真
root = Root(full_system=False, system=system)
m5.instantiate()
print(f"Starting simulation with {args.cpu_type} CPU...")
exit_event = m5.simulate()
print(f"Exiting @ tick {m5.curTick()} because {exit_event.getCause()}")
